import torch
from d2l import torch as d2l

torch.set_printoptions(2)

def multibox_prior(data,sizes,ratios):
    height,width = data.shape[-2:]
    device,num_sizes,num_radios = data.device,len(sizes),len(ratios)
    boxes_per_pixel = (num_sizes + num_radios - 1)
    size_tensor = torch.tensor(sizes,device=device)
    ratio_tensor = torch.tensor(ratios,device=device)

    offset_h,offset_w = 0.5,0.5
    steps_h = 1.0/height
    steps_w = 1.0/width

    center_h = (torch.arange(height,device=device) + offset_h) * steps_h
    center_w = (torch.arange(width,device=device) + offset_w) * steps_w

    shift_y, shift_x = torch.meshgrid(center_h,center_w)
    shift_y, shift_x = shift_y.reshape(-1),shift_x.reshape(-1)

    w = torch.cat((size_tensor * torch.sqrt(ratio_tensor[0]),
                   size_tensor[0] * torch.sqrt(ratio_tensor[1:]))) * height / width
    h = torch.cat((size_tensor / torch.sqrt(ratio_tensor[0]),
                   size_tensor[0] / torch.sqrt(ratio_tensor[1:])))
    
    anchor_manipulations = torch.stack((-w,-h,w,h)).T.repeat(height * width)
    out_grid = torch.stack([shift_x,shift_y,shift_x,shift_y],dim=1)\
                            .repeat_interleave(boxes_per_pixel)
    output = out_grid + anchor_manipulations
    return output.unsqueeze(0)